package cn.com.basic.framework.aliyun.odps.utils;

import cn.com.basic.framework.aliyun.odps.error.OdpsError;
import cn.com.basic.framework.exception.exceptions.SysException;
import cn.com.basic.framework.utils.RetryUtil;
import cn.hutool.core.collection.CollectionUtil;
import com.aliyun.odps.*;
import com.aliyun.odps.data.RecordReader;
import com.aliyun.odps.task.SQLTask;
import com.aliyun.odps.tunnel.TableTunnel;
import com.aliyun.odps.tunnel.io.ProtobufRecordPack;
import com.aliyun.odps.tunnel.io.TunnelRecordWriter;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.regex.Pattern;

/**
 * 工具类
 *
 * @author zhuxj
 */
public class OdpsUtil {

    private static Logger logger = LoggerFactory.getLogger(OdpsUtil.class);

    /**
     * 分区值校验，排除分区参数为空，是分区表但表没有分区值，传递了一个不存在的分区值的情况
     *
     * @param table     表名
     * @param partition 分区列表
     */
    public static void checkPartition(Table table, List<String> partition) {

        if (CollectionUtil.isEmpty(partition)) {
            throw SysException.newException(OdpsError.PARTITION_ERROR, String.format("分区信息没有配置.由于源头表:%s 为分区表, 所以您需要配置其抽取的表的分区信息. 格式形如:pt=hello,ds=hangzhou，请您参考此格式修改该配置项.", table.getName()));
        } else {
            List<String> allPartitions = OdpsUtil.getTableAllPartitions(table);

            if (null == allPartitions || allPartitions.isEmpty()) {
                throw SysException.newException(OdpsError.PARTITION_ERROR, String.format("分区信息配置错误.源头表:%s 虽然为分区表, 但其实际分区值并不存在. 请确认源头表已经生成该分区，再进行数据抽取.", table.getName()));
            }

            List<String> parsedPartitions = expandUserConfiguredPartition(allPartitions, partition);

            if (null == parsedPartitions || parsedPartitions.isEmpty()) {
                throw SysException.newException(OdpsError.PARTITION_ERROR, String.format("分区配置错误，根据您所配置的分区没有匹配到源头表中的分区. 源头表所有分区是:[\n%s\n], 您配置的分区是:[\n%s\n]. 请您根据实际情况在作出修改. ", StringUtils.join(allPartitions, "\n"), StringUtils.join(partition, "\n")));
            }
        }

    }

    /**
     * 获得表所有分区
     *
     * @param table 分区表
     * @return
     */
    public static List<String> getTableAllPartitions(Table table) {
        List<Partition> tableAllPartitions = table.getPartitions();

        List<String> retPartitions = new ArrayList<String>();

        if (null != tableAllPartitions) {
            for (Partition partition : tableAllPartitions) {
                retPartitions.add(partition.getPartitionSpec().toString());
            }
        }

        return retPartitions;
    }

    /**
     * 扩展用户的分区列信息
     *
     * @param allPartitions
     * @param userConfiguredPartitions
     * @return
     */
    public static List<String> expandUserConfiguredPartition(List<String> allPartitions, List<String> userConfiguredPartitions) {
        // 对odps 本身的所有分区进行特殊字符的处理
        List<String> allStandardPartitions = OdpsUtil.formatPartitions(allPartitions);

        // 对用户自身配置的所有分区进行特殊字符的处理
        List<String> allStandardUserConfiguredPartitions = OdpsUtil.formatPartitions(userConfiguredPartitions);

        /**
         *  对配置的分区级数(深度)进行检查
         *  (1)先检查用户配置的分区级数,自身级数是否相等
         *  (2)检查用户配置的分区级数是否与源头表的的分区级数一样
         */
        String firstPartition = allStandardUserConfiguredPartitions.get(0);
        int firstPartitionDepth = firstPartition.split(",").length;

        String comparedPartition;
        int comparedPartitionDepth;
        for (int i = 1, len = allStandardUserConfiguredPartitions.size(); i < len; i++) {
            comparedPartition = allStandardUserConfiguredPartitions.get(i);
            comparedPartitionDepth = comparedPartition.split(",").length;
            if (comparedPartitionDepth != firstPartitionDepth) {
                throw SysException.newException(OdpsError.PARTITION_ERROR,
                        String.format("分区配置错误, 您所配置的分区级数和该表的实际情况不一致, 比如分区:[%s] 是 %s 级分区, 而分区:[%s] 是 %s 级分区. DataX 是通过英文逗号判断您所配置的分区级数的. 正确的格式形如\"pt=0, type=0\" ，请您参考示例修改该配置项. ",
                                firstPartition, firstPartitionDepth, comparedPartition, comparedPartitionDepth));
            }
        }

        int tableOriginalPartitionDepth = allStandardPartitions.get(0).split(",").length;
        if (firstPartitionDepth != tableOriginalPartitionDepth) {
            throw SysException.newException(OdpsError.PARTITION_ERROR,
                    String.format("分区配置错误, 您所配置的分区:%s 的级数:%s 与您要读取的 ODPS 源头表的分区级数:%s 不相等. DataX 是通过英文逗号判断您所配置的分区级数的.正确的格式形如\"pt=0, type=0\" ，请您参考示例修改该配置项.",
                            firstPartition, firstPartitionDepth, tableOriginalPartitionDepth));
        }

        List<String> retPartitions = filterByRegulars(allStandardPartitions,
                allStandardUserConfiguredPartitions);

        return retPartitions;
    }

    public static String formatPartition(String partition) {
        if (StringUtils.isBlank(partition)) {
            throw SysException.newException(OdpsError.PARTITION_ERROR, "您所配置的分区不能为空白.");
        } else {
            return partition.trim().replaceAll(" *= *", "=")
                    .replaceAll(" */ *", ",").replaceAll(" *, *", ",")
                    .replaceAll("'", "");
        }
    }

    public static List<String> formatPartitions(List<String> partitions) {
        if (null == partitions || partitions.isEmpty()) {
            return Collections.emptyList();
        } else {
            List<String> formattedPartitions = new ArrayList<String>();
            for (String partition : partitions) {
                formattedPartitions.add(formatPartition(partition));

            }
            return formattedPartitions;
        }
    }

    //已经去重
    public static List<String> filterByRegular(List<String> allStrs,
                                               String regular) {
        List<String> matchedValues = new ArrayList<String>();

        // 语法习惯上的兼容处理(pt=* 实际正则应该是：pt=.*)
        String newReqular = regular.replace(".*", "*").replace("*", ".*");

        Pattern p = Pattern.compile(newReqular);

        for (String partition : allStrs) {
            if (p.matcher(partition).matches()) {
                if (!matchedValues.contains(partition)) {
                    matchedValues.add(partition);
                }
            }
        }

        return matchedValues;
    }

    //已经去重
    public static List<String> filterByRegulars(List<String> allStrs,
                                                List<String> regulars) {
        List<String> matchedValues = new ArrayList<String>();

        List<String> tempMatched = null;
        for (String regular : regulars) {
            tempMatched = filterByRegular(allStrs, regular);
            if (null != tempMatched && !tempMatched.isEmpty()) {
                for (String temp : tempMatched) {
                    if (!matchedValues.contains(temp)) {
                        matchedValues.add(temp);
                    }
                }
            }
        }

        return matchedValues;
    }

    public static TableTunnel.DownloadSession createMasterSessionForPartitionedTable(Odps odps, String tunnelServer, final String projectName, final String tableName, String partition) {

        final TableTunnel tunnel = new TableTunnel(odps);
        if (StringUtils.isNoneBlank(tunnelServer)) {
            tunnel.setEndpoint(tunnelServer);
        }

        final PartitionSpec partitionSpec = new PartitionSpec(partition);

        try {
            return RetryUtil.executeWithRetry(new Callable<TableTunnel.DownloadSession>() {
                @Override
                public TableTunnel.DownloadSession call() throws Exception {
                    return tunnel.createDownloadSession(
                            projectName, tableName, partitionSpec);
                }
            }, 10, 1000, true);
        } catch (Exception e) {
            throw SysException.newException(OdpsError.TUNNEL_ERROR, e);
        }
    }

    public static TableTunnel.DownloadSession createMasterSessionForNonPartitionedTable(Odps odps, String tunnelServer,
                                                                                        final String projectName, final String tableName) {

        final TableTunnel tunnel = new TableTunnel(odps);
        if (StringUtils.isNoneBlank(tunnelServer)) {
            tunnel.setEndpoint(tunnelServer);
        }

        try {
            return RetryUtil.executeWithRetry(new Callable<TableTunnel.DownloadSession>() {
                @Override
                public TableTunnel.DownloadSession call() throws Exception {
                    return tunnel.createDownloadSession(
                            projectName, tableName);
                }
            }, 10, 1000, true);
        } catch (Exception e) {
            throw SysException.newException(OdpsError.TUNNEL_ERROR, e);
        }
    }

    public static RecordReader getRecordReader(final TableTunnel.DownloadSession downloadSession, final long start, final long count,
                                               final boolean isCompress) {
        try {
            return RetryUtil.executeWithRetry(new Callable<RecordReader>() {
                @Override
                public RecordReader call() throws Exception {
                    return downloadSession.openRecordReader(start, count, isCompress);
                }
            }, 10, 1000, true);
        } catch (Exception e) {
            throw SysException.newException(OdpsError.TUNNEL_ERROR,
                    "open RecordReader失败. 请联系 ODPS 管理员处理." + e.getMessage());
        }
    }

    public static void preCheckPartition(Odps odps, Table table, String partition, boolean truncate) {
        boolean isPartitionedTable = OdpsUtil.isPartitionedTable(table);

        if (truncate) {
            //需要 truncate
            if (isPartitionedTable) {
                //分区表
                if (StringUtils.isBlank(partition)) {
                    throw SysException.newException(OdpsError.PARTITION_ERROR, String.format("您没有配置分区信息，因为你配置的表是分区表:%s 如果需要进行 truncate 操作，必须指定需要清空的具体分区. 请修改分区配置，格式形如 pt=${bizdate} .",
                            table.getName()));
                }
            } else {
                //非分区表
                if (StringUtils.isNotBlank(partition)) {
                    throw SysException.newException(OdpsError.PARTITION_ERROR, String.format("分区信息配置错误，你的ODPS表是非分区表:%s 进行 truncate 操作时不需要指定具体分区值. 请检查您的分区配置，删除该配置项的值.",
                            table.getName()));
                }
            }
        } else {
            //不需要 truncate
            if (isPartitionedTable) {
                //分区表
                if (StringUtils.isBlank(partition)) {
                    throw SysException.newException(OdpsError.PARTITION_ERROR,
                            String.format("您的目的表是分区表，写入分区表:%s 时必须指定具体分区值. 请修改您的分区配置信息，格式形如 格式形如 pt=${bizdate}.", table.getName()));
                }
            } else {
                //非分区表
                if (StringUtils.isNotBlank(partition)) {
                    throw SysException.newException(OdpsError.PARTITION_ERROR,
                            String.format("您的目的表是非分区表，写入非分区表:%s 时不需要指定具体分区值. 请删除分区配置信息", table.getName()));
                }
            }
        }
    }

    public static void dealTruncate(Odps odps, Table table, String partition, boolean truncate) {
        boolean isPartitionedTable = OdpsUtil.isPartitionedTable(table);

        if (truncate) {
            //需要 truncate
            if (isPartitionedTable) {
                //分区表
                if (StringUtils.isBlank(partition)) {
                    throw SysException.newException(OdpsError.PARTITION_ERROR, String.format("您没有配置分区信息，因为你配置的表是分区表:%s 如果需要进行 truncate 操作，必须指定需要清空的具体分区. 请修改分区配置，格式形如 pt=${bizdate} .",
                            table.getName()));
                } else {
                    logger.info("Try to truncate partition=[{}] in table=[{}].", partition, table.getName());
                    OdpsUtil.truncatePartition(odps, table, partition);
                }
            } else {
                //非分区表
                if (StringUtils.isNotBlank(partition)) {
                    throw SysException.newException(OdpsError.PARTITION_ERROR, String.format("分区信息配置错误，你的ODPS表是非分区表:%s 进行 truncate 操作时不需要指定具体分区值. 请检查您的分区配置，删除该配置项的值.",
                            table.getName()));
                } else {
                    logger.info("Try to truncate table:[{}].", table.getName());
                    OdpsUtil.truncateNonPartitionedTable(odps, table);
                }
            }
        } else {
            //不需要 truncate
            if (isPartitionedTable) {
                //分区表
                if (StringUtils.isBlank(partition)) {
                    throw SysException.newException(OdpsError.PARTITION_ERROR,
                            String.format("您的目的表是分区表，写入分区表:%s 时必须指定具体分区值. 请修改您的分区配置信息，格式形如 格式形如 pt=${bizdate}.", table.getName()));
                } else {
                    boolean isPartitionExists = OdpsUtil.isPartitionExist(table, partition);
                    if (!isPartitionExists) {
                        logger.info("Try to add partition:[{}] in table:[{}].",
                                partition, table.getName());
                        OdpsUtil.addPart(odps, table, partition);
                    }
                }
            } else {
                //非分区表
                if (StringUtils.isNotBlank(partition)) {
                    throw SysException.newException(OdpsError.PARTITION_ERROR,
                            String.format("您的目的表是非分区表，写入非分区表:%s 时不需要指定具体分区值. 请删除分区配置信息", table.getName()));
                }
            }
        }
    }

    public static boolean isPartitionedTable(Table table) {
        return getPartitionDepth(table) > 0;
    }

    public static int getPartitionDepth(Table table) {
        TableSchema tableSchema = table.getSchema();

        return tableSchema.getPartitionColumns().size();
    }

    public static void truncatePartition(Odps odps, Table table, String partition) {
        if (isPartitionExist(table, partition)) {
            dropPart(odps, table, partition);
        }
        addPart(odps, table, partition);
    }

    private static boolean isPartitionExist(Table table, String partition) {
        // check if exist partition 返回值不为 null
        List<String> odpsParts = OdpsUtil.listOdpsPartitions(table);

        int j = 0;
        for (; j < odpsParts.size(); j++) {
            if (odpsParts.get(j).replaceAll("'", "").equals(partition)) {
                break;
            }
        }

        return j != odpsParts.size();
    }

    public static List<String> listOdpsPartitions(Table table) {
        List<String> parts = new ArrayList<String>();
        try {
            List<Partition> partitions = table.getPartitions();
            for (Partition partition : partitions) {
                parts.add(partition.getPartitionSpec().toString());
            }
        } catch (Exception e) {
            throw SysException.newException(OdpsError.PARTITION_ERROR, String.format("获取 ODPS 目的表:%s 的所有分区失败. 请联系 ODPS 管理员处理. ",
                    table.getName()) + e.getMessage());
        }
        return parts;
    }

    private static void dropPart(Odps odps, Table table, String partition) {
        String partSpec = getPartSpec(partition);
        StringBuilder dropPart = new StringBuilder();
        dropPart.append("alter table ").append(table.getName())
                .append(" drop IF EXISTS partition(").append(partSpec)
                .append(");");
        try {
            runSqlTaskWithRetry(odps, dropPart.toString(), 10, 1000, true);
        } catch (Exception e) {
            throw SysException.newException(OdpsError.PARTITION_ERROR,
                    String.format("Drop  ODPS 目的表分区失败. 错误发生在项目:%s 的表:%s 的分区:%s .请联系 ODPS 管理员处理. ",
                            table.getProject(), table.getName(), partition) + e.getMessage());
        }
    }

    private static String getPartSpec(String partition) {
        StringBuilder partSpec = new StringBuilder();
        String[] parts = partition.split(",");
        for (int i = 0; i < parts.length; i++) {
            String part = parts[i];
            String[] kv = part.split("=");
            if (kv.length != 2) {
                throw SysException.newException(OdpsError.PARTITION_ERROR,
                        String.format("ODPS 目的表自身的 partition:%s 格式不对. 正确的格式形如: pt=1,ds=hangzhou", partition));
            }
            partSpec.append(kv[0]).append("=");
            partSpec.append("'").append(kv[1].replace("'", "")).append("'");
            if (i != parts.length - 1) {
                partSpec.append(",");
            }
        }
        return partSpec.toString();
    }

    public static void runSqlTaskWithRetry(final Odps odps, final String query, int retryTimes,
                                           long sleepTimeInMilliSecond, boolean exponential) throws Exception {
        for (int i = 0; i < retryTimes; i++) {
            try {
                runSqlTask(odps, query);
                return;
            } catch (SysException e) {
                if (OdpsError.RUN_SQL_ODPS_EXCEPTION.equals(e.getError().getKey())) {
                    logger.debug("Exception when calling callable", e);
                    if (i + 1 < retryTimes && sleepTimeInMilliSecond > 0) {
                        logger.warn(String.format("will do [%s] times retry, current exception=%s", i + 1, e.getMessage()));
                        long timeToSleep;
                        if (exponential) {
                            timeToSleep = sleepTimeInMilliSecond * (long) Math.pow(2, i);
                            if (timeToSleep >= 128 * 1000) {
                                timeToSleep = 128 * 1000;
                            }
                        } else {
                            timeToSleep = sleepTimeInMilliSecond;
                            if (timeToSleep >= 128 * 1000) {
                                timeToSleep = 128 * 1000;
                            }
                        }

                        try {
                            Thread.sleep(timeToSleep);
                        } catch (InterruptedException ignored) {
                        }
                    } else {
                        throw e;
                    }
                } else {
                    throw e;
                }
            } catch (Exception e) {
                throw e;
            }
        }
    }

    public static void runSqlTask(Odps odps, String query) {
        if (StringUtils.isBlank(query)) {
            return;
        }

        String taskName = "datax_odpswriter_trunacte_" + UUID.randomUUID().toString().replace('-', '_');

        logger.info("Try to start sqlTask:[{}] to run odps sql:[\n{}\n] .", taskName, query);

        Instance instance;
        Instance.TaskStatus status;
        try {
            instance = SQLTask.run(odps, odps.getDefaultProject(), query, taskName, null, null);
            instance.waitForSuccess();
            status = instance.getTaskStatus().get(taskName);
            if (!Instance.TaskStatus.Status.SUCCESS.equals(status.getStatus())) {
                throw SysException.newException(OdpsError.RUN_SQL_FAILED,
                        String.format("ODPS 目的表在运行 ODPS SQL失败, 返回值为:%s. 请联系 ODPS 管理员处理. SQL 内容为:[\n%s\n].", instance.getTaskResults().get(taskName),
                                query));
            }
        } catch (SysException e) {
            throw e;
        } catch (Exception e) {
            throw SysException.newException(OdpsError.RUN_SQL_ODPS_EXCEPTION,
                    String.format("ODPS 目的表在运行 ODPS SQL 时抛出异常, 请联系 ODPS 管理员处理. SQL 内容为:[\n%s\n]. ", query) + e.getMessage());
        }
    }

    public static void addPart(Odps odps, Table table, String partition) {
        String partSpec = getPartSpec(partition);
        // add if not exists partition
        StringBuilder addPart = new StringBuilder();
        addPart.append("alter table ").append(table.getName()).append(" add IF NOT EXISTS partition(")
                .append(partSpec).append(");");
        try {
            runSqlTaskWithRetry(odps, addPart.toString(), 10, 1000, true);
        } catch (Exception e) {
            throw SysException.newException(OdpsError.PARTITION_ERROR,
                    String.format("添加 ODPS 目的表的分区失败. 错误发生在添加 ODPS 的项目:%s 的表:%s 的分区:%s. 请联系 ODPS 管理员处理. ",
                            table.getProject(), table.getName(), partition) + e.getMessage());
        }
    }

    public static void truncateNonPartitionedTable(Odps odps, Table tab) {
        String truncateNonPartitionedTableSql = "truncate table " + tab.getName() + ";";

        try {
            runSqlTaskWithRetry(odps, truncateNonPartitionedTableSql, 10, 1000, true);
        } catch (Exception e) {
            throw SysException.newException(OdpsError.TABLE_TRUNCATE_ERROR,
                    String.format(" 清空 ODPS 目的表:%s 失败, 请联系 ODPS 管理员处理. ", tab.getName()) + e.getMessage());
        }
    }

    public static TableTunnel.UploadSession createMasterTunnelUpload(final TableTunnel tunnel, final String projectName,
                                                                     final String tableName, final String partition) {
        if (StringUtils.isBlank(partition)) {
            try {
                return RetryUtil.executeWithRetry(new Callable<TableTunnel.UploadSession>() {
                    @Override
                    public TableTunnel.UploadSession call() throws Exception {
                        return tunnel.createUploadSession(projectName, tableName);
                    }
                }, 10, 1000L, true);
            } catch (Exception e) {
                throw SysException.newException(OdpsError.TUNNEL_ERROR,
                        "创建TunnelUpload失败. 请联系 ODPS 管理员处理.  " + e.getMessage());
            }
        } else {
            final PartitionSpec partitionSpec = new PartitionSpec(partition);
            try {
                return RetryUtil.executeWithRetry(new Callable<TableTunnel.UploadSession>() {
                    @Override
                    public TableTunnel.UploadSession call() throws Exception {
                        return tunnel.createUploadSession(projectName, tableName, partitionSpec);
                    }
                }, 10, 1000L, true);
            } catch (Exception e) {
                throw SysException.newException(OdpsError.TUNNEL_ERROR,
                        "创建TunnelUpload失败. 请联系 ODPS 管理员处理. " + e.getMessage());
            }
        }
    }

    public static void slaveWriteOneBlock(final TableTunnel.UploadSession slaveUpload, final ProtobufRecordPack protobufRecordPack,
                                          final long blockId, final boolean isCompress) {
        try {
            RetryUtil.executeWithRetry(new Callable<Void>() {
                @Override
                public Void call() throws Exception {
                    TunnelRecordWriter tunnelRecordWriter = (TunnelRecordWriter) slaveUpload.openRecordWriter(blockId, isCompress);
                    tunnelRecordWriter.write(protobufRecordPack);
                    tunnelRecordWriter.close();
                    return null;
                }
            }, 10, 1000L, true);
        } catch (Exception e) {
            throw SysException.newException(OdpsError.ODPS_WRITER_ERROR,
                    String.format("ODPS 目的表写 block:%s 失败， uploadId=[%s]. 请联系 ODPS 管理员处理. ", blockId, slaveUpload.getId()) + e.getMessage());
        }

    }

    public static void masterCompleteBlocks(final TableTunnel.UploadSession masterUpload, final Long[] blocks) {
        try {
            RetryUtil.executeWithRetry(new Callable<Void>() {
                @Override
                public Void call() throws Exception {
                    masterUpload.commit(blocks);
                    return null;
                }
            }, 10, 1000L, true);
        } catch (Exception e) {
            throw SysException.newException(OdpsError.TUNNEL_ERROR,
                    String.format("ODPS 目的表在提交 block:[\n%s\n] 时失败, uploadId=[%s]. 请联系 ODPS 管理员处理. ", StringUtils.join(blocks, ","), masterUpload.getId()) + e.getMessage());
        }
    }
}
